# vim:ts=4:sw=4:noet:
""" Control Flow Graph """

import logging
import ir

cfgcount = 1

class CFGBasicBlock:
	def __init__(self,name=""):
		global cfgcount
		self.statements = []
		self.prevs = []
		self.nexts = []
		self.references = []
		self.name = "%d_%s" % (cfgcount, name)
		cfgcount += 1

	def addStatement(self, st):
		self.statements.append(st)

	def addNext(self, n):
		n.addPrev(self)
		self.nexts.append(n)
		if len(self.nexts)>2:
			raise Exception("%s has %d nexts: %s" % (self.name, len(self.nexts), self.nexts))

	def removeNexts(self):
		for n in self.nexts:
			n.removePrev(self)
		self.nexts = []

	def removePrev(self, p):
		self.prevs.pop(self.prevs.index(p))

	def addPrev(self, p):
		self.prevs.append(p)

	def addReference(self, r):
		self.references.append(r)

	def isEntry(self):
		return len(self.prevs)==0

	def isExit(self):
		return len(self.nexts)==0

	def cleanup(self, visited=[]):
		if self.name in visited:
			return
		visited.append(self.name)
		eraseNexts = False
		# add extra jumps if needed (for connector bbs)
		if not self.isExit():
			if (len(self.statements)>0):
				last = self.statements[-1]
				if isinstance(last, ir.IRReturn):
					eraseNexts = True
				elif not last.isJump():
					if len(self.nexts)!=1:
						raise Exception("invalid num lens: %s" % self)
					self.addStatement(CFGJump(self.nexts[0]))
			else:
				self.addStatement(CFGJump(self.nexts[0]))
		# merge connector bbs where possible
		if len(self.nexts)==1:
			if (len(self.statements)==0):
				raise Exception("problem with num statements in bb %s" % self)
			if (len(self.statements)==1) and isinstance(self.statements[0],CFGJump):
				target = self.statements[0].target
				ok = True
				for p in self.prevs:
					if len(p.nexts)>1:
						ok = False
						break
				if ok:
					for p in self.prevs:
						p.nexts = []
						p.addNext(target)
					for r in self.references:
						r.updateRef(self, target)
		# clean nexts
		for n in self.nexts:
			n.cleanup(visited)
		# erase next links if this should be an exit block
		if eraseNexts:
			self.removeNexts()
	
	def cg(self, codegen, visited=[], stopat=[]):
		if self.name in stopat:
			return
		if self.name in visited:
			return
		visited.append(self.name)
		codegen.label(self.name)
		for st in self.statements:
			st.cg(codegen)
		if (not self.isExit()) and (len(self.statements)>0):
			last = self.statements[-1]
			last.genJump(codegen, visited, stopat)
		else:
			for n in self.nexts:
				n.cg(codegen, visited, stopat)

	def dorepr(self, visited):
		if self.name in visited:
			return ""
		visited.append(self.name)
		info = ""
		if self.isEntry(): info += "entry "
		if self.isExit(): info += "exit "
		ret = "(bb %s %s\n" % (self.name, info)
		for s in self.statements:
			ret += "  %s\n" % s
		ret += "  (next %s)\n" % map(lambda x: x.name, self.nexts)
		ret += ")\n"
		for n in self.nexts:
			ret += n.dorepr(visited)
		return ret

	def __repr__(self):
		return self.dorepr([])[:-1]

class CFGIf:
	def __init__(self, tb, fb, eb, iftrue=True):
		self.tb = tb
		self.fb = fb
		self.eb = eb
		self.eb.addReference(self)		
		self.iftrue = iftrue
		
	def cg(self, codegen):
		pass

	def updateRef(self, old, new):
		self.eb = new
		self.eb.addReference(self)
		
	def isJump(self):
		return True
	
	def genJump(self, codegen, visited, stop):
		nstop = [self.eb.name]
		for s in stop: nstop.append(s)
		if self.iftrue:
			codegen.branchIfTrue(self.tb.name)
			self.fb.cg(codegen, visited, stopat=nstop)
			self.tb.cg(codegen, visited, stopat=stop)
		else:
			codegen.branchIfFalse(self.fb.name)
			self.tb.cg(codegen, visited, stopat=nstop)
			self.fb.cg(codegen, visited, stopat=stop)
		if self.eb.name not in visited and len(self.eb.prevs)>0:
			self.eb.cg(codegen, visited, stopat=stop)

	def __repr__(self):
		return "(cfgif %d)" % self.iftrue

class CFGJump:
	def __init__(self, target):
		self.target = target
		self.target.addReference(self)
	
	def cg(self, codegen):
		pass
		
	def updateRef(self, old, new):
		self.target = new
		self.target.addReference(self)

	def isJump(self):
		return True
	
	def genJump(self, codegen, visited, stop):
		if (self.target.name in visited) or (self.target.name in stop):
			codegen.branch(self.target.name)
		self.target.cg(codegen, visited, stop)

	def __repr__(self):
		return "(cfgjump)"

class CFGEndStatement:
	def cg(self, codegen):
		codegen.pop()
	
	def isJump(self):
		return False
	
	def __repr__(self):
		return "(cfgendstmt)"

class CFGEndBlock:
	def cg(self, codegen):
		codegen.returnStack()
	
	def isJump(self):
		return False
	
	def __repr__(self):
		return "(cfgendblk)"

class CFGEndBody:
	def cg(self, codegen):
		codegen.returnSelf()
	
	def isJump(self):
		return False
	
	def __repr__(self):
		return "(cfgendbody)"

class CFGStartBlock:
	def __init__(self, tempStart, blkbb, endbb):
		self.tempStart = tempStart
		self.blkbb = blkbb
		self.endbb = endbb

	def isJump(self):
		return True
	
	def genJump(self, cgen, visited, stop):
		cgen.pushBlock(self.tempStart, self.endbb.name)
		self.blkbb.cg(cgen, visited, stop)
		self.endbb.cg(cgen, visited, stop)
	
	def cg(self, cgen):
		pass

	def __repr__(self):
		return "(cfgstartblk)"

